#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import numpy as np
import re
import argparse
from pathlib import Path
import pdb

DT_F32 = 'f32'
DT_F64 = 'f64'
DT_I8 = 'int8'
DT_U8 = 'uint8'

def __build_arg_parser():
    parser = argparse.ArgumentParser(prog='inner_data_2_npy.py'
            , description='把txt格式的内部数据输出文件变成.npy文件')

    parser.add_argument('-i'
            , '--input'
            , type = str
            , help = '指定要读入的内部数据输出文件(txt格式)')

    parser.add_argument('-t'
            , '--tag'
            , type = str
            , required = False
            , help = '指定内部数据输出文件的tag，只保存tag的数据.缺省保存所有tag的数据')

    parser.add_argument('-sd'
            , '--save-dir'
            , type = str
            , required = True
            , help = '指定保存数据文件目录')

    parser.add_argument( '-dt'
            , '--data-type'
            , type = str
            , choices = [DT_F32, DT_F64, DT_I8, DT_U8]
            , default = DT_F32
            , help = f'选择使用的numpy.dtype, 缺省{DT_F32}')

    parser.add_argument('-f'
            , '--force'
            , action = 'store_true'
            , help = '如果指定的文件已经存在，覆盖之')

    return parser

def __get_dtype(dataType):
    if dataType == DT_F32:
        return np.float32
    if dataType == DT_F64:
        return np.float64
    if dataType == DT_I8:
        return np.int8
    if dataType == DT_U8:
        return np.uint8
    else:
        raise Exception(f'Unknown data type {dataType}')

def __check_save_dir(args):
    p = Path(args.save_dir)
    if p.exists() and not p.is_dir():
        raise Exception(f'{p} is not a directory')

    if not p.exists():
        p.mkdir(exist_ok=True, parents=True)
    return p

def __get_data(inputFile, tag):
    shapeKey = re.compile(f'{tag} shape:\(')
    dataKey = re.compile(f'{tag}:\[')
    with open(inputFile) as f:
        for line in f:
            msk = shapeKey.search(line)
            if msk is not None and msk.start() == 0:
                # 去掉末尾的)
                ss = line[msk.end():-2]
                shape = [int(s) for s in ss.split(',')]
            mdk = dataKey.search(line)
            if mdk is not None and mdk.start() == 0:
                # 去掉末尾的]
                ds = line[mdk.end():-2] 
                data = [float(d) for d in ds.split(',')]
    return shape, data

def __get_all_tags(inFile):
    tags = []
    with open(inFile) as f:
        for line in f:
            l = line.split(' shape:')
            if len(l) == 2:
                tags.append(l[0])
    return tags

def __convert_to_nparray(shape, data, dtype):
    a = np.array(data).astype(dtype)
    return a.reshape(shape)

def main():
    parser = __build_arg_parser()
    args = parser.parse_args()

    saveDir = __check_save_dir(args)

    if args.tag is None:
        tags = __get_all_tags(args.input)
    else:
        tags = [ args.tag ]
    print(tags)

    dType = __get_dtype(args.data_type)
    for tag in tags:
        shape, data = __get_data(args.input, tag)
        npData = __convert_to_nparray(shape, data, dType)
        f = saveDir / f'{tag}.npy'
        np.save(f, npData)
        print(f'Saved {tag} in {f}')

if __name__ == '__main__':
    try:
        main()
    except Exception as e:
        import traceback
        traceback.print_exc()
